from contextlib import asynccontextmanager
from asyncio import current_task
from functools import lru_cache
import traceback

from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine, async_scoped_session
from sqlalchemy.orm import sessionmaker
from app.config import (
    MYSQL_HOST, MYSQL_PORT, MYSQL_USER, MYSQL_PASSWORD, MYSQL_DATABASE
)

@lru_cache
def get_mysql_engine():
    return create_async_engine(
        f"mysql+aiomysql://{MYSQL_USER}:{MYSQL_PASSWORD}@{MYSQL_HOST}:{MYSQL_PORT}/{MYSQL_DATABASE}", 
        echo=False,
        pool_size=10,
        max_overflow=20,
        pool_timeout=30,
        pool_recycle=60*5,
        pool_pre_ping=True
    )

@asynccontextmanager
async def get_mysql_session():
    """
    处理session进入离开上下文
    """
    session_factory = async_scoped_session(
        sessionmaker(
            get_mysql_engine(), 
            class_=AsyncSession,
            expire_on_commit=False
        ),
        scopefunc=current_task
    )
    session:AsyncSession = session_factory()
    try:
        yield session
        await session.commit()
    except Exception as e:
        await session.rollback()
        raise
    finally:
        await session.close()
